import tushare as ts
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import os
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class DataFetcher:
    def __init__(self, token):
        """
        初始化数据获取器
        :param token: Tushare Pro API token
        """
        ts.set_token(token)
        self.pro = ts.pro_api()
        self.data_dir = 'data'
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir)
    
    def get_stock_list(self, market='A'):
        """
        获取股票列表
        :param market: 市场类型，默认为A股
        :return: 股票列表DataFrame
        """
        try:
            # 获取沪深A股列表
            stock_list = self.pro.stock_basic(exchange='', list_status='L')
            logger.info(f"成功获取{len(stock_list)}支股票信息")
            return stock_list
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return pd.DataFrame()
    
    def get_daily_data(self, ts_code, start_date, end_date):
        """
        获取股票日频OHLCV数据
        :param ts_code: 股票代码
        :param start_date: 开始日期，格式：YYYYMMDD
        :param end_date: 结束日期，格式：YYYYMMDD
        :return: 日频数据DataFrame
        """
        try:
            # 检查是否有本地缓存
            file_path = os.path.join(self.data_dir, f"{ts_code}_daily_{start_date}_{end_date}.csv")
            if os.path.exists(file_path):
                df = pd.read_csv(file_path)
                logger.info(f"从本地加载{ts_code}日频数据")
                return df
            
            # 从API获取数据
            df = ts.pro_bar(ts_code=ts_code, adj='qfq', start_date=start_date, end_date=end_date)
            if df.empty:
                logger.warning(f"{ts_code}日频数据为空")
                return df
            
            # 保存到本地
            df.to_csv(file_path, index=False)
            logger.info(f"成功获取并保存{ts_code}日频数据")
            return df
        except Exception as e:
            logger.error(f"获取{ts_code}日频数据失败: {e}")
            return pd.DataFrame()
    
    def get_financial_indicators(self, ts_code, start_date, end_date):
        """
        获取季度财务指标
        :param ts_code: 股票代码
        :param start_date: 开始日期，格式：YYYYMMDD
        :param end_date: 结束日期，格式：YYYYMMDD
        :return: 财务指标DataFrame
        """
        try:
            # 检查是否有本地缓存
            file_path = os.path.join(self.data_dir, f"{ts_code}_financial_{start_date}_{end_date}.csv")
            if os.path.exists(file_path):
                df = pd.read_csv(file_path)
                logger.info(f"从本地加载{ts_code}财务数据")
                return df
            
            # 从API获取数据
            start_year = int(start_date[:4])
            end_year = int(end_date[:4])
            
            dfs = []
            for year in range(start_year, end_year + 1):
                for quarter in range(1, 5):
                    # 获取单季度财务指标
                    df_quarter = self.pro.fina_indicator(ts_code=ts_code, period=f"{year}{quarter:02}01")
                    dfs.append(df_quarter)
            
            if not dfs:
                logger.warning(f"{ts_code}财务数据为空")
                return pd.DataFrame()
            
            df = pd.concat(dfs, ignore_index=True)
            
            # 保存到本地
            df.to_csv(file_path, index=False)
            logger.info(f"成功获取并保存{ts_code}财务数据")
            return df
        except Exception as e:
            logger.error(f"获取{ts_code}财务数据失败: {e}")
            return pd.DataFrame()
    
    def get_hs300_data(self, start_date, end_date):
        """
        获取沪深300指数数据
        :param start_date: 开始日期，格式：YYYYMMDD
        :param end_date: 结束日期，格式：YYYYMMDD
        :return: 沪深300指数数据DataFrame
        """
        try:
            # 检查是否有本地缓存
            file_path = os.path.join(self.data_dir, f"hs300_{start_date}_{end_date}.csv")
            if os.path.exists(file_path):
                df = pd.read_csv(file_path)
                logger.info(f"从本地加载沪深300指数数据")
                return df
            
            # 获取沪深300指数数据
            df = ts.pro_bar(ts_code='399300.SZ', asset='I', start_date=start_date, end_date=end_date)
            
            # 保存到本地
            df.to_csv(file_path, index=False)
            logger.info(f"成功获取并保存沪深300指数数据")
            return df
        except Exception as e:
            logger.error(f"获取沪深300指数数据失败: {e}")
            return pd.DataFrame()

# 使用示例
if __name__ == "__main__":
    # 请替换为你的Tushare Pro token
    token = "your_tushare_token"
    fetcher = DataFetcher(token)
    
    # 获取股票列表
    stocks = fetcher.get_stock_list()
    print(f"共获取到{len(stocks)}支股票")
    
    # 获取单支股票日频数据
    if not stocks.empty:
        ts_code = stocks.iloc[0]['ts_code']
        daily_data = fetcher.get_daily_data(ts_code, '20200101', '20231231')
        print(f"{ts_code}日频数据: {len(daily_data)}条记录")
        
        # 获取财务数据
        financial_data = fetcher.get_financial_indicators(ts_code, '20200101', '20231231')
        print(f"{ts_code}财务数据: {len(financial_data)}条记录")
    
    # 获取沪深300指数数据
    hs300_data = fetcher.get_hs300_data('20200101', '20231231')
    print(f"沪深300指数数据: {len(hs300_data)}条记录")    